package cn.hy.limiter;

import java.lang.reflect.Method;
import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;

/**
 * 对于很多情况下至少需要支持按照IP和方法名，或者其它自定义的方式进行限流。
 * 单机限流
 *
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspect {

    //private final ConcurrentHashMap<String, RateLimiter> EXISTED_RATE_LIMITERS = new ConcurrentHashMap<>();

    // 防止内存撑爆
    private final Cache<String,RateLimiter> EXISTED_RATE_LIMITERS = CacheBuilder.newBuilder()
            .maximumSize(10000)
            .expireAfterAccess(Duration.ofHours(1))
            .build();

    @Pointcut("@annotation(cn.hy.limiter.RateLimit)")
    public void rateLimit() {
    }

    @Around("rateLimit()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        String canonicalName = point.getTarget().getClass().getCanonicalName();
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        RateLimit annotation = AnnotationUtils.findAnnotation(method, RateLimit.class);
        String cacheKey = canonicalName+"#"+method.getName();
        // get rate limiter
        // RateLimiter rateLimiter = EXISTED_RATE_LIMITERS.computeIfAbsent(cacheKey, k -> RateLimiter.create(annotation.limit()));

        RateLimiter rateLimiter = EXISTED_RATE_LIMITERS.get(cacheKey, () ->  RateLimiter.create(annotation.limit()));
        // process
        if (rateLimiter!=null && rateLimiter.tryAcquire()) {
            return point.proceed();
        } else {
            throw new RuntimeException("too many requests, please try again later...");
        }
    }
}